

from major.models_manager import chat_model
from langgraph.prebuilt import create_react_agent
from typing import Annotated
import datetime

def current_time():
    """返回时间的工具"""
    print("正在获取时间...")
    return datetime.datetime.now()

def add_numbers(
    a: Annotated[float, "第一个加数"],
    b: Annotated[float, "第二个加数"]
) -> float:
    """执行两个数的加法运算"""
    print("正在执行加法运算...")
    return a + b


model = chat_model.get_model()
# 创建Agent，包含两个工具：calculator 和 current_time
agent = create_react_agent(model,
                           tools=[current_time,add_numbers],
                           prompt='''你是一个助手''')


response  = agent.astream_events({'messages': [{'role': 'user', 'content': '现在几点'}]})
#print(response['messages'][-1].content)


async def stream_response():
    async for event in response:
        kind = event["event"]
        if kind == 'on_chat_model_stream':
            content = event["data"]["chunk"].content
            if content:
                # 空内容意味着模型正在请求调用一个工具。因此，我们只打印非空内容。”
                print(content, end="|", flush=True)

# 运行异步主函数
if __name__ == "__main__":
    import asyncio
    asyncio.run(stream_response())